#from gtk._gtk import Orientable
from scipy import ndimage
PREFIX_U = 'u_r_'
PREFIX_V = 'v_r_'
PREFIX_RHO = 'rho_r_'
PATH = '/home/behollis/DATA/in/ts00050/'
DPATH = '/home/behollis/DATA/out/lockExSt/ts00050/'
#ODPATH = '/home/behollis/Dropbox/eurovis2015shortpaperCode/oslines/'
OUTPUT_DATA_DIR = '/home/behollis/Dropbox/eurovis2015shortpaperCode/oslines/'

import math
import netCDF4
import math as pm
import sys
import csv
import loadStreamlines
import h5py
from loadStreamlines import readStreamlines
import interpProbDensity
import matplotlib.pyplot as plt
import numpy as np
import matplotlib
from netcdf_reader import *
import scipy


FILE_NAME = 'pe_dif_sep2_98.nc' 
FILE_NAME_CENTRAL_FORECAST = 'pe_fct_aug25_sep2.nc'
INPUT_DATA_DIR = '/home/behollis/DATA/pierre/ocean/'
OUTPUT_DATA_DIR = '/home/behollis/Dropbox/eurovis2015shortpaperCode/oslines/'
COM =  2
LON = 53
LAT = 90
LEV = 16
DEPTH = -2.0
INTEGRATION_DIR = 'a'

#prob = None
#slines = None

#from http://stackoverflow.com/questions/8661537/how-to-perform-bilinear-interpolation-in-python
def bilinear_interpolation(x, y, points):
    '''Interpolate (x,y) from values associated with four points.

    The four points are a list of four triplets:  (x, y, value).
    The four points can be in any order.  They should form a rectangle.

        >>> bilinear_interpolation(12, 5.5,
        ...                        [(10, 4, 100),
        ...                         (20, 4, 200),
        ...                         (10, 6, 150),
        ...                         (20, 6, 300)])
        165.0

    '''
    # See formula at:  http://en.wikipedia.org/wiki/Bilinear_interpolation
    
    # 

    points = sorted(points)               # order points by x, then by y
    (x1, y1, q11), (_x1, y2, q12), (x2, _y1, q21), (_x2, _y2, q22) = points

    if x1 != _x1 or x2 != _x2 or y1 != _y1 or y2 != _y2:
        raise ValueError('points do not form a rectangle')
    if not x1 <= x <= x2 or not y1 <= y <= y2:
        raise ValueError('(x, y) not within the rectangle')

    return (q11 * (x2 - x) * (y2 - y) +
            q21 * (x - x1) * (y2 - y) +
            q12 * (x2 - x) * (y - y1) +
            q22 * (x - x1) * (y - y1)
           ) / ((x2 - x1) * (y2 - y1) + 0.0)

def interpInCell(ppos=[0.0,0.0],member = 0):
    #decompose fract / whole from particle position
    ppos_parts = [[0.0,0.0],[0.0,0.0]] #[fract,whole] for each x,y comp
    ppos_parts[0][0] = pm.modf(ppos[0])[0];ppos_parts[0][1] = pm.modf(ppos[0])[1]
    ppos_parts[1][0] = pm.modf(ppos[1])[0];ppos_parts[1][1] = pm.modf(ppos[1])[1]
    
    #print "ensemble alpha x: " + str( ppos_parts[0][0] )
    #print "ensemble alpha y: " + str( ppos_parts[1][0] )
    
    # grid point numbers:
    #
    # (2)---(3)
    # |      |
    # |      |
    # (0)---(1)
    
    #find four corner grid point indices, numbered from gpt0 = (bottom, left) TO gpt3 = (top, right)
    #calculated from whole parts 
    gpt0 = [ppos_parts[0][1], ppos_parts[1][1]]
    gpt1 = [ppos_parts[0][1] + 1, ppos_parts[1][1]]
    gpt2 = [ppos_parts[0][1], ppos_parts[1][1] + 1]
    gpt3 = [ppos_parts[0][1] + 1, ppos_parts[1][1] + 1]
    
    
    #vclin[curr_realization][curr_lat][curr_lon][curr_level][0] 
    #vclin[curr_realization][curr_lat][curr_lon][curr_level][1] 
    
    try:
        gpt0_velx = vclin[member][gpt0[0]][gpt0[1]][SEED_LEVEL][0]
        gpt0_vely = vclin[member][gpt0[0]][gpt0[1]][SEED_LEVEL][1]
         
        gpt1_velx = vclin[member][gpt1[0]][gpt1[1]][SEED_LEVEL][0]
        gpt1_vely = vclin[member][gpt1[0]][gpt1[1]][SEED_LEVEL][1] 
        
        gpt2_velx = vclin[member][gpt2[0]][gpt2[1]][SEED_LEVEL][0]
        gpt2_vely = vclin[member][gpt2[0]][gpt2[1]][SEED_LEVEL][1] 
        
        gpt3_velx = vclin[member][gpt3[0]][gpt3[1]][SEED_LEVEL][0]
        gpt3_vely = vclin[member][gpt3[0]][gpt3[1]][SEED_LEVEL][1] 
    except:
        return 0.0, 0.0
    
    lerp_velx = bilinear_interpolation(ppos[0], ppos[1], [(gpt0[0],gpt0[1],gpt0_velx),(gpt1[0],gpt1[1],gpt1_velx),\
                                                          (gpt2[0],gpt2[1],gpt2_velx),(gpt3[0],gpt3[1],gpt3_velx)])
    lerp_vely = bilinear_interpolation(ppos[0], ppos[1], [(gpt0[0],gpt0[1],gpt0_vely),(gpt1[0],gpt1[1],gpt1_vely),\
                                                          (gpt2[0],gpt2[1],gpt2_vely),(gpt3[0],gpt3[1],gpt3_vely)])
    
    return lerp_velx, lerp_vely
            
def readVelNetCDF(file):
    # open a the netCDF file for reading.
    ncfile = netCDF4.Dataset(file,'r') 
    # read the data in variable named 'data'.
    data = ncfile.variables['ks_test_stat'][:]
    #nx,ny = data.shape
    # check the data.
    #data_check = arange(nx*ny) # 1d array
    #data_check.shape = (nx,ny) # reshape to 2d array

    # close the file.
    ncfile.close()
    
    return data 


def calcStreamlineProbs2(seed_pt, sl_dict, MEM):
    
    slines_prob = dict()
    
    # gather all of the vector fields
    #print len(sl_dict.keys())
    
    for midx in sl_dict.keys():
        #print midx
        #vfield = interpProbDensity.readMember( int(midx[3:]) )
        interpProbDensity.gUv[midx] = vclin[int(midx),:,:,SEED_LEVEL,:]
        interpProbDensity.gUv[midx] = np.swapaxes(interpProbDensity.gUv[midx],1,2)
        interpProbDensity.gUv[midx] = np.swapaxes(interpProbDensity.gUv[midx],0,1)
        
    
    probs_list = list()#contains all prob densities
    
#    for slidx, sl in enumerate(slines):
    for slkey in sl_dict.keys():
        sl = sl_dict[slkey]
        #tot = 0
        slines_prob[slkey] = [ 0.0, [] ]
        for ptidx in range(0, len(sl[0])):
            
            #print 'streamline length: ' + str( len(sl[0]) )
            #print 'for sl key: ' + str(slkey)
            
            pt = [ sl[0][ptidx], sl[1][ptidx] ]
            
            try: 
                kde, min_u, max_u, min_v, max_v = interpProbDensity.interpVelFromEnsemble(pt[0:2], MEM)
                
                #plotDistro( kde, (min_u,max_u), (min_v,max_v), '') 
                #band_width = MEM**(-1./(DIM+4))
                EPS_FACTOR = 0.1
                epsx = EPS_FACTOR * math.fabs(max_u - min_u)
                epsy = EPS_FACTOR * math.fabs(max_v - min_v)
                
                u = interpProbDensity.gUv[slkey][0][pt[0],pt[1]]
                v = interpProbDensity.gUv[slkey][1][pt[0],pt[1]]
                low_bounds = [u - epsx, v - epsy]
                high_bounds = [u + epsx, v + epsy]
                prob = kde.integrate_box( low_bounds, high_bounds )
                
                if prob < 0.0:
                    prob = 0
                
                #SCALE = 100
                #prob *= SCALE
                
               
                slines_prob[slkey][0] += prob  
                
                #print prob
                #print slines_prob[slkey][0]
                
                #tot += prob #Possibly make this a Markov chain (however densities are very low)
                    
                #print 'for pt: ' + str(ptidx)
                #print 'prob' + str(prob)
                #print 'tot' + str(tot)
                    
                slines_prob[slkey][1].append(float(prob))
                
                
                probs_list.append(prob)
                
                    
                #print('streamline {0}: {1}'.format(slkey, tot))
            except:
                print 'streamlines integrated out of sim domain'
                continue
            
        
    return slines_prob, probs_list

from operator import itemgetter
import random

g_crisp_streamlines = list()
TOTAL_STEPS = 550#int(sys.argv[6])
INTEGRATION_DIR = 'a' #str(sys.argv[7]).lower()
SEED_LAT = 42#float(sys.argv[2])
SEED_LON = 30#float(sys.argv[3])
SEED_LEVEL = 0#int(sys.argv[4])
integration_step_size = 0.01#float(sys.argv[5])

def readStreamlinePositions(data, filename):
    #change to 'wb' after initial debug...
    filename = OUTPUT_DATA_DIR + filename
    reader = csv.reader(open(filename + ".csv", 'r'), delimiter=',')
    idx = 0
    for row in reader:
        #print row
        data[idx] = [float(i) for i in row]
        idx += 1 
        
def readParticles():
    str_integration_values = '_ss' + str(integration_step_size) + '_ts' + str(TOTAL_STEPS) + '_dir_' + str(INTEGRATION_DIR)

    #read crisp sphaghetti plots
    for idx in range(0,MEM):
        curr_member_sl = [[],[],[]]
        readStreamlinePositions(curr_member_sl,'lat'+str(SEED_LAT)+'_lon'+str(SEED_LON)+'_lev'+str(SEED_LEVEL)+'_mem'+str(idx)+str_integration_values) 
        #g_crisp_streamlines.append(curr_member_sl)
        sl_dict[str(idx)] = curr_member_sl
        
sl_dict = dict()
        
        
def renderOverallLikelihoods(top=None):
    from numpy import linalg
    
    fig = plt.figure()
    plt.box('off')
    plt.imshow(im)
    
    plt.tick_params(\
    axis='both',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',
    right='off',
    left='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off', labelleft='off') # labels along the bottom edge are off
    
    # Render with likelihoods
    for tup in toplist:
        slkey = tup[0]
#        if top is not None and slkey in top:
        sl = sl_dict[slkey]
        tot = tup[1]
        clr = mapper2.to_rgba(x=tot) 
        
        sl = np.asarray(sl)
        
        alpha = tot / max2 #normalize magnitude for alpha                             
        flip = np.zeros( len(sl[1]) ); flip.fill( im.shape[1] )
        
        plt.plot(IMGSCALEX*(LAT-sl[0] + OFFSET), IMGSCALEY*(LON-sl[1]), \
                 color=(clr[0], clr[1], clr[2], alpha), linewidth = 1.0 )
        
#            plt.plot( IMGSCALEX*np.asarray( sl[0] ), \
#                     flip - IMGSCALEY*np.asarray(sl[1]),\
#                     color=(clr[0], clr[1], clr[2], alpha), linewidth = 1.0 )
      
    plt.plot([ IMGSCALEX*(LAT-SEED_LAT + OFFSET) ], [ IMGSCALEY*(LON-SEED_LON) ], '+', color='black',linewidth = LW, ms=MS ) # draw seed loc 
#    plt.xlim([xmin, xmax])
#    plt.ylim([ymin, ymax])  
    #plt.colorbar(mapper2)   
    #plt.show() 
    plt.savefig('../img/topProb.png')    

def renderIndividualSlines():
    # Rendering streamlines by velocity prob density per integration step
    
    '''
    http://matplotlib.org/examples/pylab_examples/multicolored_line.html
    http://stackoverflow.com/questions/19390895/matplotlib-plot-with-variable-line-width
    '''
    
    from matplotlib.collections import LineCollection
    
    for slkey in sl_dict.keys():
        
        fig = plt.figure()
        #ax = fig.gca(projection='3d')
        axes = fig.add_subplot(1, 1, 1, axisbg='white')
        axes.set_aspect('equal', 'datalim')
        # render streamlines
        
        sl = sl_dict[slkey]
       
        try:
            l = list(slines_prob[slkey][1])
            #print 'length of list l: ' + str( len(l) )
            
            min_= min(l)
            max_ = max(l)
        except:
            #print 'streamline not more than one integration step'
            continue
        
        # draw each segment a color depending on prob density
        steps = len( sl[0] ) - 1
        
        x = list()
        y = list()
        lws = list()
        clrs = list()
        
        for idx in range(0, steps):
           
            #normalize by shifting / scaling to [0,1]
            scale = (slines_prob[slkey][1][idx] - min_0) / math.fabs( max_0 - min_0) 
        
                
            prob = slines_prob[slkey][1][idx]
            print prob
            '''
            if scale < 0:
                scale = 0.0
            print scale
            '''
            clr = mapper.to_rgba(x=prob) 
            print clr  
            #clr = (1.0*scale, 1.0*scale, 1.0*scale, 1.0)
            x0 = sl[0][idx]; x1 = sl[0][idx+1]
            y0 = sl[1][idx]; y1 = sl[1][idx+1]
            SCALE_W = 3.0
            MIN_W = 1.0
            #axes.plot([x0, x1 ], [y0, y1], color=clr, linewidth = scale*SCALE_W + MIN_W )
            x.append(LAT-x0); x.append(LAT-x1)
            y.append(y0), y.append(y1)
            lws.append(scale*SCALE_W + MIN_W)
            clrs.append(clr)
            
        points = np.array([x,y]).T.reshape(-1,1,2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        lc = LineCollection(segments, linewidths=lws, colors=clrs)
        axes.add_collection(lc)
        axes.plot([ sx ], [ sy ], '+', color='black', linewidth=LW , ms=MS) # draw seed loc
#        plt.plot([ IMGSCALEX*(LAT-SEED_LAT + OFFSET) ], [ IMGSCALEY*(LON-SEED_LON) ], '+', color='black',linewidth = LW, ms=MS ) # draw seed loc 
            
        #plt.colorbar(mapper)
        
        #plt.xlim(x0=50)
        #plt.ylim(x0=50)
        #plt.xlim([50,100])
        #plt.ylim([50,100])
        plt.savefig('../img/sline{0}.png'.format(slkey))

def renderColorbars():
    #get images of colorbars
    # make a color map of fixed colors
    fig, ax = plt.subplots()
    fig.colorbar(mapper, orientation='horizontal', ticks=[min_0, max_0] )
    plt.savefig('../img/colorbar1.png')
    fig, ax = plt.subplots()
    fig.colorbar(mapper2, orientation='vertical', ticks=[min2, max2])
    plt.savefig('../img/colorbar2.png')
    
def get_line(x1, y1, x2, y2):
    '''
    http://www.roguebasin.com/index.php?title=Bresenham%27s_Line_Algorithm
    http://www.cs.helsinki.fi/group/goa/mallinnus/lines/bresenh.html
    '''
    
    points = []
    issteep = abs(y2-y1) > abs(x2-x1)
    if issteep:
        x1, y1 = y1, x1
        x2, y2 = y2, x2
    rev = False
    if x1 > x2:
        x1, x2 = x2, x1
        y1, y2 = y2, y1
        rev = True
    deltax = x2 - x1
    deltay = abs(y2-y1)
    error = int(deltax / 2)
    y = y1
    ystep = None
    if y1 < y2:
        ystep = 1
    else:
        ystep = -1
    for x in range(x1, x2 + 1):
        if issteep:
            points.append((y, x))
        else:
            points.append((x, y))
        error -= deltay
        if error < 0:
            y += ystep
            error += deltax
    # Reverse the list if the coordinates were reversed
    if rev:
        points.reverse()
    return points

def computeHistogramEqual(buf):
    '''
    http://stackoverflow.com/questions/5858902/histogram-equalization-of-matplotlib-color-tables
    '''
    '''
    Non-linear colormap...
    http://stackoverflow.com/questions/22521382/nonlinear-colormap-matplotlib
    '''
    divs = 1000
    imvals = np.sort(buf.flatten())
    lo = imvals[0]
    hi = imvals[-1]
    steps = (imvals[::len(imvals)/divs] - lo) / (hi - lo)
    num_steps = float(len(steps))
    interps = [(s, idx/num_steps, idx/num_steps) for idx, s in enumerate(steps)]
    interps.append((1, 1, 1))
    cdict = {'red' : interps,
             'green' : interps,
             'blue' : interps}
    histeq_cmap = matplotlib.colors.LinearSegmentedColormap('HistEq', cdict)
    
    return histeq_cmap

class nlcmap(object):
    '''
    http://stackoverflow.com/questions/22521382/nonlinear-colormap-matplotlib
    '''
    def __init__(self, cmap, levels):
        self.cmap = cmap
        self.N = cmap.N
        self.monochrome = self.cmap.monochrome
        self.levels = np.asarray(levels, dtype='float64')
        self._x = self.levels
        self.levmax = self.levels.max()
        self.transformed_levels = np.linspace(0.0, self.levmax,
             len(self.levels))

    def __call__(self, xi, alpha=1.0, **kw):
        yi = np.interp(xi, self._x, self.transformed_levels)
        return self.cmap(yi / self.levmax, alpha)
    
def renderMultipleSlineSegments(toplist=None):
    
    '''
    plot line segments probability (additive) into buffer using
    http://en.wikipedia.org/wiki/Bresenham%27s_line_algorithm
    http://www.roguebasin.com/index.php?title=Bresenham%27s_Line_Algorithm
    '''
    
    buf = np.zeros( shape=im.shape )
    
    all_points = list()
    
    for slkey in sl_dict.keys():    
        sl = sl_dict[slkey]
        
        if toplist is not None and slkey in toplist:
            steps = len( sl[0] ) - 1
            for idx in range(0, steps):   
                #normalize by shifting / scaling to [0,1]
                prob = slines_prob[slkey][1][idx]    
                
                x0 = sl[0][idx]; x1 = sl[0][idx+1]
                y0 = sl[1][idx]; y1 = sl[1][idx+1]
                
                # scale to input LIC image buffer size
                x0sc = IMGSCALEX * x0
                y0sc = im.shape[1] - IMGSCALEY*y0
                x1sc = IMGSCALEX * x1
                y1sc = im.shape[1] - IMGSCALEY*y1
                
        #            points = get_line(int(x0),int(y0),int(x1),int(y1))
                points = get_line(int(round(x0sc)),\
                                  int(round(y0sc)),\
                                  int(round(x1sc)),\
                                  int(round(y1sc)))
                
                # add prob vel density to buffer
                for pt in points:
                    all_points.append(pt)
                    curr_prob = buf[pt[1], pt[0],0]
                    if prob > curr_prob:
                        buf[pt[1], pt[0],0] += prob
                    #print buf[pt[1], pt[0],0]
                
   
    #http://stackoverflow.com/questions/17201172/a-logarithmic-colorbar-in-matplotlib-scatter-plot
    '''
    The way that matplotlib does color mapping is in two steps, first a Normalize function 
    (wrapped up by the sub-classes of matplotlib.colors.Normalize) which maps the data you 
    hand in to [0, 1]. The second step maps values in [0,1] -> RGBA space.
    You just need to use the LogNorm normalization class, passed in with the norm kwarg.
    '''
    '''
    map = matplotlib.cm.ScalarMappable(norm=matplotlib.colors.LogNorm())
    map.set_array( buf[:,:,0] )
    map.set_cmap('gist_rainbow')#sequential
    map.autoscale()
    min, max = map.get_clim()
    '''
    min = np.amin(buf[:,:,0])
    max = np.amax(buf[:,:,0])
    
#    hmap = computeHistogramEqual(buf)

    rg = max-min
    
    #the choice of the levels depends on the data:
    levels = np.concatenate((
        np.linspace(max - 0.1*rg, max, 3),np.linspace(min, min + 0.1*rg, 3)
        ))
    
    levels = levels[levels <= max]
    levels.sort()
    
    cmap_nonlin = nlcmap(plt.cm.jet, levels)
    
    #for the colorbar we map the original colormap, not the nonlinear one:
    sm = plt.cm.ScalarMappable(cmap=plt.cm.jet, 
                    norm=plt.Normalize(vmin=min, vmax=max))
    sm._A = []
    
#    cbar = fig.colorbar(sm, cax=cbar_ax)
    #here we are relabel the linear colorbar ticks to match the nonlinear ticks
#    cbar.set_ticks(cmap_nonlin.transformed_levels)
#    cbar.set_ticklabels(["%.2f" % lev for lev in levels])
    
    fig, ax = plt.subplots()
    cbar = fig.colorbar(sm, orientation='horizontal')#, ticks=[min, max])
    cbar.set_ticks(cmap_nonlin.transformed_levels)
    cbar.set_ticklabels(["%.4f" % lev for lev in levels])
    plt.savefig('../img/colorbarCumulativeTest.png')
    
    #fig, ax = plt.subplots()
    #fig.colorbar(hmap, orientation='horizontal')#, ticks=[min, max] )
    #plt.savefig('../img/colorbarHistoCmap.png')
    
    maxp = np.amax(buf[:,:,0])
#    minp = np.amin(buf)
    
    for pt in all_points:
        #print x
        #print y
        #if buf[x,y,0] < 0.0001:
        #    buf[x,y,:] = im[x,y,:]
        #else:
        x = pt[1]; y = pt[0]
        #http://graph-plotter.cours-de-math.eu/
        #plotting this response for alpha
        #alpha = math.log( (buf[x,y,0] / maxp) + 1.0 ) + 0.3 
#        buf[x,y,:] = map.to_rgba(buf[x,y,0])
        buf[x,y,:] = cmap_nonlin(buf[x,y,0])       
#        buf[x,y,:] = map.to_rgba(hmap(buf[x,y,0])[0])
        #buf[x,y,3] = alpha

    # draw seed loc
    plt.figure() 
    plt.box('off')
    plt.imshow(im)
    imgplot = plt.imshow(buf) 
    
    plt.plot([ IMGSCALEX*sx ], [ im.shape[1] - IMGSCALEY*sy ], '+', color='black',linewidth = LW, ms=MS )
    
    plt.tick_params(\
    axis='both',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',
    right='off',
    left='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off', labelleft='off') # labels along the bottom edge are off
    
    plt.xlim([xmin, xmax])
    plt.ylim([ymin, ymax])
    
    plt.savefig('../img/cumulativeProbsSlinesTest.png')
    #plt.show()
    
def renderMultipleSlineSegments2(toplist=None):
    
    '''
    plot line segments probability (additive) into buffer using
    http://en.wikipedia.org/wiki/Bresenham%27s_line_algorithm
    http://www.roguebasin.com/index.php?title=Bresenham%27s_Line_Algorithm
    '''
    
    buf = np.zeros( shape=im.shape )
    
    all_points = list()
    
    for slkey in sl_dict.keys():    
        sl = sl_dict[slkey]
        
        if toplist is not None and slkey in toplist:
            steps = len( sl[0] ) - 1
            for idx in range(0, steps):   
                #normalize by shifting / scaling to [0,1]
                prob = slines_prob[slkey][1][idx]    
                
                x0 = sl[0][idx]; x1 = sl[0][idx+1]
                y0 = sl[1][idx]; y1 = sl[1][idx+1]
                
                # scale to input LIC image buffer size
                x0sc = IMGSCALEX * x0
                y0sc = im.shape[1] - IMGSCALEY*y0
                x1sc = IMGSCALEX * x1
                y1sc = im.shape[1] - IMGSCALEY*y1
                
                rr, cc, val = line_aa(int(round(y0sc)),\
                                  int(round(x0sc)),\
                                  int(round(y1sc)),\
                                  int(round(x1sc)))
                
                for idx in range(0, len(val)):
                    all_points.append((rr[idx], cc[idx]))
                    if buf[rr[idx], cc[idx], 0] < prob * val[idx]:
                        buf[rr[idx], cc[idx], :] = [prob * val[idx]]*4
                    
    min = np.amin(buf[:,:,0])
    max = np.amax(buf[:,:,0])

    rg = max-min
    
    #the choice of the levels depends on the data:
    levels = np.concatenate((
        np.linspace(max - 0.1*rg, max, 3),np.linspace(min, min + 0.1*rg, 3)
        ))
    
    levels = levels[levels <= max]
    levels.sort()
    
    cmap_nonlin = nlcmap(plt.cm.jet, levels)
    
    #for the colorbar we map the original colormap, not the nonlinear one:
    sm = plt.cm.ScalarMappable(cmap=plt.cm.jet, 
                    norm=plt.Normalize(vmin=min, vmax=max))
    sm._A = []
    
    fig, ax = plt.subplots()
    cbar = fig.colorbar(sm, orientation='horizontal')#, ticks=[min, max])
    cbar.set_ticks(cmap_nonlin.transformed_levels)
    cbar.set_ticklabels(["%.4f" % lev for lev in levels])
    plt.savefig('../img/colorbarCumulativeTestAA.png')
   
    maxp = np.amax(buf[:,:,0])

    for pt in all_points:
        x = pt[0]; y = pt[1]
        buf[x,y,:] = cmap_nonlin(buf[x,y,0])       
                
    # draw seed loc
    plt.figure() 
    plt.box('off')
    plt.imshow(im)
    imgplot = plt.imshow(buf) 
    
    plt.plot([ IMGSCALEX*sx ], [ im.shape[1] - IMGSCALEY*sy ], '+', color='black',linewidth = LW, ms=MS )
    
    plt.tick_params(\
    axis='both',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',
    right='off',
    left='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off', labelleft='off') # labels along the bottom edge are off
    
    plt.xlim([xmin, xmax])
    plt.ylim([ymin, ymax])
    
    plt.savefig('../img/cumulativeProbsSlinesTestAA.png')
    #plt.show()
    
def renderMultipleSlineSegments3(top=None):
    from matplotlib.collections import LineCollection
    
    toplist_dict = dict(toplist)
    
    fig = plt.figure()
    #ax = fig.gca(projection='3d')
    axes = fig.add_subplot(1, 1, 1, axisbg='white')
    #axes.set_aspect('equal', 'datalim')
    # render streamlines
    
    plt.box('off')
    plt.imshow(im)
   
    for slkey in sl_dict.keys():
        if top is not None and slkey in top:
            sl = sl_dict[slkey]
           
            try:
                l = list(slines_prob[slkey][1])
                #print 'length of list l: ' + str( len(l) )
                
                min_= min(l)
                max_ = max(l)
            except:
                #print 'streamline not more than one integration step'
                continue
            
            # draw each segment a color depending on prob density
            steps = len( sl[0] ) - 1
            
            x = list()
            y = list()
            lws = list()
            clrs = list()
            
            for idx in range(0, steps):
               
                #normalize by shifting / scaling to [0,1]
                scale = (slines_prob[slkey][1][idx] - min_0) / math.fabs( max_0 - min_0) 
            
                    
                prob = slines_prob[slkey][1][idx]
                print prob
                '''
                if scale < 0:
                    scale = 0.0
                print scale
                '''
                clr = list(mapper.to_rgba(x=prob)) 
                tot = toplist_dict[slkey]
                alpha = tot / max2 #normalize magnitude for alpha  
                clr[3] = alpha
                print clr  
                #clr = (1.0*scale, 1.0*scale, 1.0*scale, 1.0)
                x0 = sl[0][idx]; x1 = sl[0][idx+1]
                y0 = sl[1][idx]; y1 = sl[1][idx+1]
                
                # scale to input LIC image buffer size
                x0sc = IMGSCALEX*(LAT-x0 + OFFSET)#IMGSCALEX * x0
                y0sc = IMGSCALEY*(LON-y0)#im.shape[1] - IMGSCALEY*y0
                x1sc = IMGSCALEX*(LAT-x0 + OFFSET)#IMGSCALEX * x1
                y1sc = IMGSCALEY*(LON-y1)#im.shape[1] - IMGSCALEY*y1
                
               
                
                SCALE_W = 3.0
                MIN_W = 1.0
                x.append(x0sc); x.append(x1sc)
                y.append(y0sc), y.append(y1sc)
                lws.append(scale*SCALE_W + MIN_W)
                
                clrs.append(tuple(clr))
                
        points = np.array([x,y]).T.reshape(-1,1,2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        lc = LineCollection(segments, linewidths=lws, colors=clrs)
        axes.add_collection(lc)
    
    plt.plot([ IMGSCALEX*(LAT-SEED_LAT + OFFSET) ], [ IMGSCALEY*(LON-SEED_LON) ], '+', color='black',linewidth = LW, ms=MS ) # draw seed loc 
 
    
    plt.tick_params(\
    axis='both',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',
    right='off',
    left='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off', labelleft='off') # labels along the bottom edge are off
    
 #   plt.xlim([xmin, xmax])
 #   plt.ylim([ymin, ymax])
    
    plt.savefig('../img/cumulativeProbsSlinesTest3.png')
    #plt.show()
   
OFFSET = 13 
def renderStandardPlot(top=None):
    ''' 
    Rank each streamline based on total KDE prob density #
    Render all for comparison
    '''
    
    fig = plt.figure(frameon=False)
    plt.box('off')
    plt.imshow(im)
    
    plt.tick_params(\
    axis='both',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',
    right='off',
    left='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off', labelleft='off') # labels along the bottom edge are off
    
    '''
    ax = fig.gca(projection='3d')
    axes = fig.add_subplot(1, 1, 1, axisbg='white')
    axes.set_aspect('equal', 'datalim')
    '''
    # render streamlines
    for slkey in sl_dict.keys():
        sl = sl_dict[slkey]
        if top is not None and slkey in top:
            sl = np.asarray(sl)
    #        flip = np.zeros( len(sl[1]) ); flip.fill( im.shape[1] )
    #            plt.plot( IMGSCALEX*np.asarray( sl[0] ), \
    #                     flip - IMGSCALEY*np.asarray(sl[1]),\
    #                     color=(0.8, 0.0, 0.0), linewidth = 1.0 )
            plt.plot(IMGSCALEX*(LAT-sl[0] + OFFSET), IMGSCALEY*(LON-sl[1]), color=(0.8, 0.0, 0.0), linewidth = 1.0)
         
    # draw seed loc  
    plt.plot([ IMGSCALEX*(LAT-SEED_LAT + OFFSET) ], [ IMGSCALEY*(LON-SEED_LON) ], '+', color='black',linewidth = LW, ms=MS ) # draw seed loc 
#    plt.plot([ IMGSCALEX*sx ], [ IMGSCALEY*sy ], '+', color='black',linewidth = LW, ms=MS )  
#    plt.xlim([0, LAT])
#    plt.ylim([0, LON])
    #plt.show()
    #lt.gcf().set_size_inches(,size/float(dpi)))
    if top is None:
        plt.savefig('../img/all.png')
    else:
        plt.savefig('../img/allTop.png')
        
if __name__ == '__main__':
    global prob, slines, MEM
    
    MEM = 600 
    PERCENT = 0.01
    readParticles()
    
    sx = LAT-SEED_LAT # seed x coord
    sy = SEED_LON # seed y coord
    
    # register plots
    XSETm = 0#10
    XSETp = 0#40
    YSETm = 0#6
    YSETp = 0#28
    
    im = plt.imread('/home/behollis/Dropbox/eurovis2015shortpaperCode/massBayPics/massBayBasecampGreyPaperImage.png')
    IMGSCALEX = (im.shape[1]) / LAT 
    IMGSCALEY = im.shape[0] / LON
    
    xmin = IMGSCALEX*sx - IMGSCALEX*XSETm
    xmax = IMGSCALEX*sx + IMGSCALEX*XSETp
    ymin = im.shape[1] - (IMGSCALEY*sy - IMGSCALEY*YSETm)
    ymax = im.shape[1] - (IMGSCALEY*sy + IMGSCALEY*YSETp)
    
    #realizations file 
    pe_dif_sep2_98_file = INPUT_DATA_DIR + FILE_NAME
    pe_fct_aug25_sep2_file = INPUT_DATA_DIR + FILE_NAME_CENTRAL_FORECAST 
    
    #realizations reader 
    rreader = NetcdfReader(pe_dif_sep2_98_file)
    
    #central forecasts reader 
    creader = NetcdfReader(pe_fct_aug25_sep2_file)
    vclin8 = creader.readVarArray('vclin', 7)
    #landmask = creader.readVarArray('landv')
    #landmask = scipy.ndimage.rotate(landmask, -90)
    
    #plt.imshow(landmask, origin='lower')
    #plt.show()
    
    #deviations from central forecast for all 600 realizations
    vclin = rreader.readVarArray('vclin')  
    vclin = addCentralForecast(vclin, vclin8, level_start=SEED_LEVEL, level_end=SEED_LEVEL)  
    
    slines_prob, prob_list = calcStreamlineProbs2([SEED_LAT, SEED_LON], sl_dict, MEM)
    
    mapper = matplotlib.cm.ScalarMappable()
    mapper2 = matplotlib.cm.ScalarMappable() #mapper for total prob streamlines
    min_0 = min(prob_list)
    max_0 = max(prob_list)
        
    shifted_prob = np.array(prob_list) #- min_0
    
    mapper.set_array( shifted_prob )
    mapper.set_cmap('jet')#sequential
    mapper.autoscale()
    
   
    
    #seed point
    LW = 10.0
    MS = 10.0
    
     # sort slines based on overall likelihood
    toplist = list()
    total_probs = list()
    for slk in sl_dict.keys():
        print slk
        toplist.append( (slk,slines_prob[slk][0]) )
        
    for key in slines_prob.keys(): 
        total_probs.append( slines_prob[key][0] )
    
    sorted( toplist, key=itemgetter(1) )
    
    toplist2 = toplist[0:int(len(toplist)*PERCENT)]
    topslines = list()
    for item in toplist2:
        topslines.append(item[0])
        
    #put most probable on top by reversing draw order with most prob last
    toplist2.reverse()
    topslines.reverse()
    
    mapper2 = matplotlib.cm.ScalarMappable() #mapper for total prob streamlines
    min2 = min(total_probs)
    max2 = max(total_probs)
    mapper2.set_array( np.array(total_probs) )
    mapper2.set_cmap('OrRd')#uential
    
    #norm = matplotlib.colors.Normalize(vmin=min_0, vmax=max_0)
    norm2 = matplotlib.colors.Normalize(vmin=min2, vmax=max2)
    mapper2.set_norm(norm2)
    #mapper.set_norm(norm)
    #mapper.autoscale()
    mapper2.autoscale()
    
    renderStandardPlot()
    renderStandardPlot(topslines)
    renderColorbars()
    renderOverallLikelihoods()
    renderIndividualSlines()
    renderMultipleSlineSegments3(topslines)
    
    print 'finished!'
        